package com.huaweicloud.obs.flinktest;

import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeutils.base.LongSerializer;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.util.Preconditions;

import java.io.Serializable;
import java.util.ArrayDeque;
import java.util.Deque;

/**
 * 功能描述
 *
 * @since 2021-05-12
 */
public class EndSequenceGenerator implements Serializable {

    private static final long serialVersionUID = 3885103821540774586L;

    private final long from;

    private final long to;

    private transient ListState<Long> state;

    protected transient Deque<Long> values;

    public EndSequenceGenerator(long from, long to) {
        this.from = from;
        this.to = to;
        this.values = new ArrayDeque<>();
    }

    public void init(String name, FunctionInitializationContext context, RuntimeContext runtime)
        throws Exception {
        Preconditions.checkState(this.state == null,
            "The " + getClass().getSimpleName() + " has already been initialized.");

        this.state = context.getOperatorStateStore().getListState(
                    new ListStateDescriptor<>(name + "-sequence-state", LongSerializer.INSTANCE));
        this.values = new ArrayDeque<>();
        if (context.isRestored()) {
            for (Long v : this.state.get()) {
                this.values.add(v);
            }
        } else {
            final int step = runtime.getNumberOfParallelSubtasks();
            final int taskIndex = runtime.getIndexOfThisSubtask();
            final long space = from + taskIndex;

            long total = Math.abs(to - from + 1);
            final int base = div(total, step);
            final int all = (total % step > taskIndex) ? base + 1 : base;

            for (long i = 0; i < all; i++) {
                this.values.add(i * step + space);
            }
        }
    }

    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        Preconditions.checkState(this.state != null,
            "The " + getClass().getSimpleName() + " state has not been properly initialized.");

        this.state.clear();
        for (Long value : this.values) {
            this.state.add(value);
        }
    }

    public boolean hasNext() {
        return !this.values.isEmpty();
    }

    public String next() {
        return values.poll().toString();
    }

    private static int div(long l, long r) {
        Preconditions.checkArgument(l >= 0);
        Preconditions.checkArgument(r > 0);
        Preconditions.checkArgument(l <= Integer.MAX_VALUE * r);
        return (int) (l / r);
    }
}
